#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import argparse
import subprocess
import shutil
import tempfile
import requests
import yaml
import re

DEFAULT_BASE_URL = "https://raw.githubusercontent.com/ros/rosdistro/master/rosdep"
BASE_YAML_NAME = "base.yaml"
PYTHON_YAML_NAME = "python.yaml"

###############################################################################
# Utilities
###############################################################################
def step_print(msg):
    """在日志中显式输出大写的步骤提示"""
    print(f"\n=== {msg} ===")

def parse_args():
    p = argparse.ArgumentParser(
        description="以自定义发行版 + 指定参考发行版进行 rosdep 自动映射（第一次+二次）并生成 YAML"
    )
    p.add_argument("--distro", required=True,
                   help="目标发行版 (例如 openeuler)")
    p.add_argument("--ref-distro", required=True,
                   help="参考发行版 (例如 fedora, debian, ubuntu 等)")
    p.add_argument("--pkg-manager", default=None,
                   help="包管理器 (默认根据 distro 猜测, 常用=dnf/apt)")
    p.add_argument("--rosdistro-path", default=None,
                   help="自定义 base.yaml / python.yaml 路径或 URL, 否则用github master")
    p.add_argument("--output-dir", default=os.getcwd(),
                   help="输出 YAML 文件夹")
    p.add_argument("--verbose", action="store_true",
                   help="输出详细失败原因报告 + 输出每个 key 查找过程")
    return p.parse_args()

def fetch_rosdep_files(rosdistro_path):
    tmpdir = tempfile.mkdtemp(prefix="rosdep_gen_")
    base_local = os.path.join(tmpdir, BASE_YAML_NAME)
    python_local = os.path.join(tmpdir, PYTHON_YAML_NAME)

    def fetch_or_copy(src, dst):
        if src.startswith("http://") or src.startswith("https://"):
            print(f"[INFO] 下载文件: {src} => {dst}")
            resp = requests.get(src, timeout=30)
            resp.raise_for_status()
            with open(dst, "wb") as f:
                f.write(resp.content)
        else:
            print(f"[INFO] 复制文件: {src} => {dst}")
            if not os.path.isfile(src):
                raise FileNotFoundError(f"{src}")
            shutil.copyfile(src, dst)

    if rosdistro_path:
        base_src = os.path.join(rosdistro_path, BASE_YAML_NAME)
        py_src   = os.path.join(rosdistro_path, PYTHON_YAML_NAME)
    else:
        base_src = f"{DEFAULT_BASE_URL}/{BASE_YAML_NAME}"
        py_src   = f"{DEFAULT_BASE_URL}/{PYTHON_YAML_NAME}"

    print(f"[INFO] 获取 base.yaml => {base_local}")
    fetch_or_copy(base_src, base_local)
    print(f"[INFO] 获取 python.yaml => {python_local}")
    fetch_or_copy(py_src, python_local)

    return base_local, python_local, tmpdir

def guess_pkg_manager(distro):
    d = distro.lower()
    if d in ("openeuler","fedora","rhel","centos","rocky","alma"):
        return "dnf"
    elif d in ("debian","ubuntu","openkylin"):
        return "apt"
    return None

def gather_local_packages(pkg_manager):
    pkg_set = set()
    if pkg_manager == "dnf":
        cmd = ["dnf", "list", "--all"]
    elif pkg_manager == "apt":
        cmd = ["apt", "list"]
    else:
        raise ValueError(f"脚本暂不支持 pkg_manager={pkg_manager}")

    print(f"[INFO] 调用 {cmd} 列出可用包...")
    try:
        res = subprocess.run(cmd, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE, text=True, check=True)
        for line in res.stdout.splitlines():
            line = line.strip()
            if not line or line.startswith("Listing") \
               or line.startswith("正在") or line.startswith("Available"):
                continue

            if pkg_manager == "apt":
                pkg = line.split('/')[0]
                pkg_set.add(pkg)
            else:
                first = line.split()[0]
                splitted = first.split('.')
                if splitted:
                    pkg_set.add(splitted[0])
    except subprocess.CalledProcessError as e:
        print(f"[ERROR] {cmd} 执行失败: {e}")
        sys.exit(1)

    print(f"[INFO] 本地 {pkg_manager} 包数量: {len(pkg_set)}")
    return pkg_set

def load_yaml(fp):
    with open(fp, "r", encoding="utf-8") as f:
        data = yaml.safe_load(f)
    return data if data else {}

def ensure_list(obj):
    if isinstance(obj, list):
        return obj
    elif isinstance(obj, str):
        return [obj]
    return []

###############################################################################
# 修正重点：对 python fallback 的 key 做“简单清理”
###############################################################################
def sanitize_python_key(key):
    """
    去除常见的前缀后缀, 让 fallback pip 包名更合理
      - 若结尾是 '-pip'，删掉
      - 若开头是 'python3-'，删掉
      (可根据需要再加其他规则)
    """
    if key.endswith("-pip"):
        key = key[: -len("-pip")]
    if key.startswith("python3-"):
        key = key[len("python3-"):]
    return key

###############################################################################
# 第一次映射
###############################################################################
def first_pass_mapping(data, ref_distro, local_pkg_set,
                       pkg_manager, is_python=False, verbose=False):
    if is_python:
        print(f"\n=== 第一次映射: python - 参考 {ref_distro} ===")
    else:
        print(f"\n=== 第一次映射: base - 参考 {ref_distro} ===")

    mapped_ok = {}
    mapped_fail = {}

    for key, val in data.items():
        if verbose:
            print(f"[INFO] 正在进行第一次映射 key: {key} (ref={ref_distro})")

        if not isinstance(val, dict):
            mapped_fail[key] = "rosdep entry is not dict"
            continue
        if ref_distro not in val:
            mapped_fail[key] = f"ref-distro={ref_distro} no mapping"
            continue

        dist_val = val[ref_distro]

        if isinstance(dist_val, dict):
            if not is_python:
                mapped_fail[key] = "pip mapping found in base.yaml?? ignoring"
                continue
            if "pip" in dist_val and isinstance(dist_val["pip"], dict):
                pkgs = dist_val["pip"].get("packages", [])
                if not pkgs:
                    mapped_fail[key] = "ref distro pip but no packages"
                else:
                    mapped_ok[key] = {"installer":"pip", "packages": pkgs}
            else:
                mapped_fail[key] = "unknown ref distro dict structure"
        elif isinstance(dist_val, list):
            pkgs = dist_val
            found_any = any(p in local_pkg_set for p in pkgs)
            if found_any:
                mapped_ok[key] = {"type":"system","pkgs": pkgs}
            else:
                mapped_fail[key] = f"system pkg not found: {pkgs}"
        elif isinstance(dist_val, str):
            if dist_val in local_pkg_set:
                mapped_ok[key] = {"type":"system","pkgs":[dist_val]}
            else:
                mapped_fail[key] = f"system pkg not found: {dist_val}"
        else:
            mapped_fail[key] = "unknown ref distro data type"

    return mapped_ok,mapped_fail

###############################################################################
# 二次查找
###############################################################################
def dnf_search_for_package(key, verbose=False):
    if verbose:
        print(f"[INFO] dnf search {key}")
    try:
        cmd = ["dnf","search", key]
        res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                             text=True, check=False)
        lines=res.stdout.splitlines()
        for line in lines:
            line=line.strip()
            if not line or line.startswith("===="):
                continue
            first=line.split()[0]
            splitted=first.split('.')
            if splitted:
                pkg_name=splitted[0]
                return pkg_name
    except Exception as e:
        if verbose:
            print(f"[ERROR] dnf search exception: {e}")
    return None

def second_pass_search(fail_dict, data, my_distro, local_pkg_set,
                       pkg_manager, is_python=False,
                       verbose=False):
    if is_python:
        print(f"\n=== 第二次查找: python (对第一次失败的进行搜索) ===")
    else:
        print(f"\n=== 第二次查找: base (对第一次失败的进行搜索) ===")

    success_map={}
    still_fail={}

    for key, reason in fail_dict.items():
        if verbose:
            print(f"[INFO] 二次查找 key: {key}, 失败原因: {reason}")
        if not is_python:
            # base => dnf search
            pkg_found = dnf_search_for_package(key, verbose=verbose)
            if pkg_found:
                success_map[key] = {"type":"system","pkgs":[pkg_found]}
            else:
                still_fail[key] = "dnf search no result"
        else:
            # python => fallback => pip
            # 修正: 对key做sanitize
            pkg_name = sanitize_python_key(key)
            success_map[key] = {"installer":"pip", "packages":[pkg_name]}

    return success_map, still_fail

###############################################################################
# 写出结果 YAML
###############################################################################
def write_output_yaml(distro, base_success_map, py_success_map, outdir):
    step_print("写出最终 YAML 文件")

    base_data={}
    for k,v in base_success_map.items():
        if "type" in v and v["type"]=="system":
            base_data[k]={ distro: v["pkgs"] }

    python_data={}
    for k,v in py_success_map.items():
        if "type" in v and v["type"]=="system":
            python_data[k]={ distro: v["pkgs"] }
        elif "installer" in v and v["installer"]=="pip":
            python_data[k]={
                distro: {
                    "pip": {
                        "packages": v["packages"]
                    }
                }
            }

    os.makedirs(outdir, exist_ok=True)
    base_path=os.path.join(outdir,f"base-{distro}.yaml")
    py_path=os.path.join(outdir,f"python-{distro}.yaml")

    with open(base_path,"w",encoding="utf-8") as f:
        yaml.dump(base_data,f,sort_keys=True)
    with open(py_path,"w",encoding="utf-8") as f:
        yaml.dump(python_data,f,sort_keys=True)

    return base_path, py_path

###############################################################################
# main
###############################################################################
def main():
    args = parse_args()

    step_print("STEP 1: 参数解析 & 猜测包管理器")
    if not args.pkg_manager:
        guess = guess_pkg_manager(args.distro)
        if guess:
            args.pkg_manager = guess
            print(f"[INFO] 未指定 --pkg-manager, 根据 {args.distro} 推断为 {args.pkg_manager}")
        else:
            print("[ERROR] 无法推断包管理器，请手动指定 --pkg-manager=xxx")
            sys.exit(1)

    step_print("STEP 2: 获取 base.yaml 与 python.yaml")
    base_fp, python_fp, tmpdir = fetch_rosdep_files(args.rosdistro_path)
    base_data   = load_yaml(base_fp)
    python_data = load_yaml(python_fp)

    step_print("STEP 3: 列出本地可用包")
    local_pkg_set = gather_local_packages(args.pkg_manager)

    step_print("STEP 4: 第一次映射 (base.yaml)")
    base_ok_1, base_fail_1 = first_pass_mapping(
        base_data, args.ref_distro, local_pkg_set,
        args.pkg_manager, is_python=False, verbose=args.verbose
    )

    step_print("STEP 5: 第一次映射 (python.yaml)")
    py_ok_1, py_fail_1 = first_pass_mapping(
        python_data, args.ref_distro, local_pkg_set,
        args.pkg_manager, is_python=True, verbose=args.verbose
    )

    step_print("STEP 6: 第二次查找 (base.yaml)")
    base_ok_2, base_fail_2 = second_pass_search(
        base_fail_1, base_data, args.distro,
        local_pkg_set, args.pkg_manager,
        is_python=False, verbose=args.verbose
    )

    step_print("STEP 7: 第二次查找 (python.yaml)")
    py_ok_2, py_fail_2 = second_pass_search(
        py_fail_1, python_data, args.distro,
        local_pkg_set, args.pkg_manager,
        is_python=True, verbose=args.verbose
    )

    step_print("STEP 8: 合并映射结果 & 统计")
    base_final = {}
    base_final.update(base_ok_1)
    base_final.update(base_ok_2)

    python_final = {}
    python_final.update(py_ok_1)
    python_final.update(py_ok_2)

    base_success = len(base_final)
    base_fail    = len(base_fail_2)
    py_success   = len(python_final)
    py_fail      = len(py_fail_2)

    print("\n=== [RESULT] ===")
    print(f"base.yaml => success: {base_success}, fail: {base_fail}")
    print(f"python.yaml => success: {py_success}, fail: {py_fail}")

    if args.verbose and (base_fail > 0 or py_fail > 0):
        print("\n[FAILURE detail] base.yaml:")
        for k,v in base_fail_2.items():
            print(f"  {k} => {v}")
        print("\n[FAILURE detail] python.yaml:")
        for k,v in py_fail_2.items():
            print(f"  {k} => {v}")

    step_print("STEP 9: 写出 YAML")
    base_path, py_path = write_output_yaml(
        args.distro, base_final, python_final, args.output_dir
    )
    print(f"[INFO] 已生成:\n  {base_path}\n  {py_path}")

    # 不更新 rosdep
    shutil.rmtree(tmpdir, ignore_errors=True)
    print("[INFO] 脚本执行完毕.")

if __name__ == "__main__":
    main()

